import matplotlib.pyplot as plt
import numpy as np
from matplotlib.ticker import MaxNLocator, MultipleLocator
import sys
from os.path import abspath, dirname
import pickle

plt.rc('font', family='Arial', weight='normal')

# 准备数据
with open("data/dcqcn_AI.txt") as f:
    datalines = f.readlines()

x_original = ['1712', '5084', '8385', '14937', '22480', '27515', '35094', '45283', '56200', '67667', '85750', '158898', '410893', '812932', '1271956', '1764734', '2776783', '4253697', '6801351', '15835374']

imgs = [[],[],[],[]]
for i in range(len(imgs)):
    for j in range(6):
        imgs[i].append([])

for i in range(0, len(datalines), 20):
    for j in range(20):
        promote = list(map(lambda x: float(x[0:-1]),datalines[i+j].split()))
        imgs[0][i // 20].append(promote[0])
        imgs[1][i // 20].append(promote[1])
        imgs[2][i // 20].append(promote[2])
        imgs[3][i // 20].append(promote[3])
x_uniform = np.arange(0,20)
xticks = []
for size in x_original:
    size = int(size)
    if size < 4000:
        xticks.append(str(int(size)) + "B")
    elif size / 1e3 < 4000:
        xticks.append(str(int(size/1e3)) + "KB")
    elif size / 1e6 < 4000:
        xticks.append(str(int(size/1e6)) + "MB")
    else:
        xticks.appedn(str(int(size/1e9)) + "GB")

# 设置通用绘图的样式
plt.figure(figsize=(12, 8))
plt.tick_params(labelsize=40)
font1 = {'weight': 'normal', 'size': 40}
X_name = "Flow Size"
Y_name = "FCT Reduction(%)"

colors = ['#1f77b4', '#1f77b4', '#ff7f0e', '#ff7f0e', '#2ca02c', '#2ca02c', '#e377c2', '#7f7f7f', '#bcbd22']
markers = ['o', 's', 'o', 's', 'o', 's', '>', 'p', '*']


# 绘制曲线
files = ["avg", "50", "95", "99"]
labels = ["30% 1xAI", "30% 5xAI", "60% 1xAI", "60% 5xAI", "90% 1xAI", "90% 5xAI"]
linestyle = ['--', '-']
for fi, fct in enumerate(imgs):
    plt.axhline(y=0, color='k', linestyle='-.', linewidth=5, alpha=0.7)
    for i, y in enumerate(fct):
        plt.plot(x_uniform, y, marker=markers[i], markersize=10, label=labels[i], linestyle=linestyle[i % len(linestyle)], linewidth=4, color=colors[i % len(colors)])

    # 添加图例
    plt.legend(loc='upper left', ncol=1, prop={'size': 20, 'weight': 'normal'}, columnspacing=1)
    plt.xlabel(X_name, fontdict=font1)
    plt.ylabel(Y_name, fontdict=font1)

    plt.gca().yaxis.set_major_locator(MaxNLocator(integer=True, prune='lower', nbins=6))  # 控制显示更多的 y 刻度
    
    plt.xticks(x_uniform, labels=xticks, rotation=45, fontsize=20)  # 设置倾斜角度和字体大小

    # 设置 x 和 y 轴的刻度样式
    plt.tick_params(which='major', axis='x', direction='in', length=10, width=4)
    plt.tick_params(which='major', axis='y', direction='in', length=10, width=4)
    plt.tick_params(which='minor', axis='x', direction='in', length=4, width=4)
    plt.tick_params(which='minor', axis='y', direction='in', length=4, width=4)

    # 保存图像
    png_filename = f'img/AI_{files[fi]}.pdf'
    plt.savefig(png_filename, format='pdf', dpi=300, bbox_inches='tight')
    png_filename = f'img/AI_{files[fi]}.png'
    plt.savefig(png_filename, format='png', dpi=300, bbox_inches='tight')

    print(f"Finish draw AI_{files[fi]}.pdf")
    plt.cla()
